package com.ai.nsa.utils.hbase;

import com.ai.nsa.utils.hbase.model.HBaseRow;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.*;

/**
 * @author Feng Chen 2017/6/7.
 */
public class HBaseOperations {

    private HBasePool hbasePool;

    private static final Logger logger = LoggerFactory.getLogger(HBaseOperations.class);

    public HBaseOperations(HBasePool hbasePool) {
        this.hbasePool = hbasePool;
    }

    public void put(String table, final String rowKey, final String family, final String qualifier, final String value) {
        put(table, new HBaseRow(rowKey, family, qualifier, value));
    }

    public void put(String table, final HBaseRow row) {
        put(table, Collections.singletonList(row));
    }

    public void put(String table, final Collection<HBaseRow> rows) {
        logger.info("Execute put request: {}", rows);
        execute(table, new Executor<Void>() {

            @Override
            public Void execute(Table table) throws Exception {
                List<Put> puts = new ArrayList<Put>(rows.size());
                for (HBaseRow row : rows) {
                    Put put = HBaseRowUtil.convert(row);
                    if (put != null) {
                        puts.add(put);
                    }
                }
                table.put(puts);
                return null;
            }
        });
    }

    public HBaseRow get(String table, final String rowKey) {
        return get(table, rowKey, new RowMapper<HBaseRow>() {
            @Override
            public HBaseRow mapRow(Result result) throws Exception {
                return HBaseRowUtil.convert(result);
            }
        });
    }

    public <T> T get(String table, final String rowKey, RowMapper<T> mapper) {
        return get(table, rowKey, null, null, mapper);
    }

    public <T> T get(String table, final String rowKey, final String family, RowMapper<T> mapper) {
        return get(table, rowKey, family, null, mapper);
    }

    public <T> T get(String table, final String rowKey, final String familyName, final String qualifier, final RowMapper<T> mapper) {

        return execute(table, new Executor<T>() {
            @Override
            public T execute(Table table) throws Exception {
                Get get = new Get(rowKey.getBytes());
                if (familyName != null) {
                    byte[] family = familyName.getBytes();
                    if (qualifier != null) {
                        get.addColumn(family, qualifier.getBytes());
                    } else {
                        get.addFamily(family);
                    }
                }
                Result result = table.get(get);
                return mapper.mapRow(result);
            }
        });
    }

    public List<HBaseRow> scan(String tableName, Scan scan) {

        return scan(tableName, scan, new ResultExtractor<List<HBaseRow>>() {

            @Override
            public List<HBaseRow> extractData(ResultScanner result) throws Exception {

                List<HBaseRow> rows = new ArrayList<HBaseRow>();
                Iterator<Result> its = result.iterator();

                while (its.hasNext()) {
                    Result row = its.next();
                    rows.add(HBaseRowUtil.convert(row));
                }
                return rows;
            }
        });
    }

    public <T> T scan(String tableName, String family, String column, final ResultExtractor<T> extractor) {
        Scan scan = new Scan();
        scan.addColumn(family.getBytes(), column.getBytes());
        return scan(tableName, scan, extractor);
    }

    public <T> T scan(String tableName, final Scan scan, final ResultExtractor<T> extractor) {
        return execute(tableName, new Executor<T>() {
            @Override
            public T execute(Table table) throws Exception {

                ResultScanner scanner = table.getScanner(scan);
                return extractor.extractData(scanner);
            }
        });
    }

    // SECTION: INNER HELPER

    private <T> T execute(String tableName, Executor<T> executor) {
        Connection connection = null;
        Table table = null;
        try {
            connection = hbasePool.borrowConnection();
            table = connection.getTable(TableName.valueOf(tableName));
            return executor.execute(table);
        } catch (Exception e) {
            logger.error("Error during HBASE operation", e);
            throw new HBaseOperationException(e);
        } finally {
            if (table != null) {
                try {
                    table.close();
                } catch (IOException e) {

                }
            }
            hbasePool.returnConnection(connection);
        }
    }


    ///////////////////////////////////////////////////////////

    public interface Executor<T> {
        T execute(Table table) throws Exception;
    }

    public interface RowMapper<T> {
        T mapRow(Result result) throws Exception;
    }

    public interface ResultExtractor<T> {
        T extractData(ResultScanner result) throws Exception;
    }


}
